"""
API路由定义，处理各种API请求
"""
from fastapi import APIRouter, HTTPException
from ..models.inference import ModelInference
from .models import (
    TextGenerationRequest,
    TextGenerationResponse,
    BatchGenerationRequest,
    BatchGenerationResponse,
    ModelInfoResponse
)
from ..config import (
    MODEL_PATH,
    MODEL_DEVICE,
    MODEL_PRECISION,
    MODEL_MAX_LENGTH,
    MODEL_MAX_NEW_TOKENS
)

# 创建路由器
router = APIRouter(prefix="/api/v1", tags=["大模型API"])

# 创建模型推理实例
inference = ModelInference()

@router.post("/generate", response_model=TextGenerationResponse)
async def generate_text(request: TextGenerationRequest):
    """
    生成文本API
    """
    try:
        result = inference.generate_text(
            prompt=request.prompt,
            max_new_tokens=request.max_new_tokens,
            temperature=request.temperature,
            top_p=request.top_p,
            top_k=request.top_k
        )
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"生成文本失败: {str(e)}")

@router.post("/batch-generate", response_model=BatchGenerationResponse)
async def batch_generate(request: BatchGenerationRequest):
    """
    批量生成文本API
    """
    try:
        results = inference.batch_generate(
            prompts=request.prompts,
            max_new_tokens=request.max_new_tokens,
            temperature=request.temperature,
            top_p=request.top_p,
            top_k=request.top_k
        )
        return {"results": results}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"批量生成文本失败: {str(e)}")

@router.get("/model-info", response_model=ModelInfoResponse)
async def get_model_info():
    """
    获取模型信息API
    """
    try:
        return {
            "model_name": MODEL_PATH.split("/")[-1],
            "model_path": MODEL_PATH,
            "device": MODEL_DEVICE,
            "precision": MODEL_PRECISION,
            "max_length": MODEL_MAX_LENGTH,
            "max_new_tokens": MODEL_MAX_NEW_TOKENS
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取模型信息失败: {str(e)}")